from classes.image_data import ImageData
from classes.position import Position
import os
import logging


def gen_origin_index(args):
    assert args.origin_dir is not None
    assert args.origin_index is not None, "Give me the filename of the created index."
    from utils.clean_text import work
    args.img_path = args.origin_dir
    args.format = 'jpg'
    args.output = args.origin_index
    work(args)


def gen_folded_img(args):
    assert args.left_up_la is not None
    assert args.right_down_la is not None
    assert args.left_up_lo is not None
    assert args.right_down_lo is not None
    assert args.origin_index is not None
    assert args.folded_index is not None, "Give me the filename of the created index."
    assert args.folded_dir is not None, "Give me the dir of the created index."
    assert args.folded_num is not None

    if os.path.exists(args.folded_dir) is False:
        os.makedirs(args.folded_dir)

    i = ImageData(args.origin_index,
                  Position(la=args.left_up_la, lo=args.left_up_lo),
                  Position(la=args.right_down_la, lo=args.right_down_lo))

    i.create_folded_img(args.folded_dir, args.folded_index, args.folded_num)


def __build_map(args):
    """Build map for no folding or folded dataset"""
    assert args.left_up_la is not None
    assert args.right_down_la is not None
    assert args.left_up_lo is not None
    assert args.right_down_lo is not None
    assert args.origin_index is not None
    if args.folded_index is None:
        # no folding
        i = ImageData(args.origin_index,
                      Position(la=args.left_up_la, lo=args.left_up_lo),
                      Position(la=args.right_down_la, lo=args.right_down_lo))
        i.build_map()
    else:
        # with folding
        assert args.folded_num is not None
        i = ImageData(args.origin_index,
                      Position(la=args.left_up_la, lo=args.left_up_lo),
                      Position(la=args.right_down_la, lo=args.right_down_lo))
        i.build_map_with_exist_fold(args.folded_index, args.folded_num)
    return i


def test_map_build(args):
    __build_map(args)
    logging.info("Test Finished.")


def test_pos_img(args):
    assert args.test_la is not None
    assert args.test_lo is not None
    i = __build_map(args)

    image_info = i.get(Position(la=args.test_la, lo=args.test_lo))
    print('Location Pic: ' + image_info.file_path)
    print('Wgs Position: ' + 'Left_up: ' + str(image_info.left_up) + ' Right_down: ' + str(image_info.right_down))


def gen_mxnet_index(args):
    assert args.mxnet_index is not None
    assert args.y_index is not None
    i = __build_map(args)
    i.ensemble_y_classification(args.y_index, args.y_scale)
    i.create_mxnet_index(args.mxnet_index, args.balance_factor, args.balance_label)
    logging.info('Building finished. Please check: ' + args.mxnet_index)


def gen_tensorflow_imgs(args):
    assert args.tensorflow_dir is not None
    assert args.y_index is not None
    i = __build_map(args)
    i.ensemble_y_classification(args.y_index, args.y_scale)
    i.create_tensorflow_imgs(args.tensorflow_dir, args.balance_factor, args.balance_label, args.tensorflow_train_factor)
    with open(os.path.join(args.tensorflow_dir, 'scale.txt'), 'w') as f:
        f.write(str(args.y_scale))
    logging.info('Building finished. Please check: ' + args.tensorflow_dir)

def gen_tensorflow_data(args):
    args.train_directory = os.path.join(args.tensorflow_dir, 'train')
    args.validation_directory = os.path.join(args.tensorflow_dir, 'validation')
    args.output_directory = args.tensorflow_dir
    args.train_shards = args.tensorflow_train_shards
    args.validation_shards = args.tensorflow_validation_shards
    args.num_threads = 2
    args.labels_file = os.path.join(args.tensorflow_dir, 'label.txt')
    from utils.tfrecord import main
    main.FLAGS = args
    main(args)
